// SPDX-License-Identifier: GPL-2.0+
/*
 * upgrade.c -- upgrade command
 *
 * Copyright (C) 2016 Toradex
 * Author: Hao Jiang <jhleiren@qq.com>
 */
#include <common.h>
#include <blk.h>
#include <command.h>
#include <console.h>
#include <cli.h>
#include <errno.h>
#include <malloc.h>

#include <upgrade.h>

extern ulong image_load_addr;	///< 这是通过 net 下载后，镜像存放在内存中的地址

static char *image_name = NULL;

/**
 * @brief 
 * 
 * @param image_type_name 
 * @return upgrade_image_type_en 
 */
static UPGRADE_IMAGE_TYPE_EN check_upgrade_image_type(char *image_type_name)
{
    upgrade_image_type_en image_type = IMAGE_TYPE_IS_UNKOWN;

    if (strcmp(image_type_name, "uboot")) {
        image_type = IMAGE_TYPE_IS_UBOOT;
    }
    else if (strcmp(image_type_name, "kernel")) {
        image_type = IMAGE_TYPE_IS_KERNEL;
    }
    else if (strcmp(image_type_name, "dts")) {
        image_type = IMAGE_TYPE_IS_DTB;
    }
    else if (strcmp(image_type_name, "rootfs")) {
        image_type = IMAGE_TYPE_IS_ROOTFS;
    }
    else (strcmp(image_type_name, "image")) {
        image_type = IMAGE_TYPE_IS_ALL;
    }

    return image_type;
}

/**
 * @brief 初始化要下载的镜像名
 * 
 * @param image_type 镜像类型
 * @param image_name 镜像名称
 * @return int 成功返回 0,失败返回 -1
 */
static char *check_upgrade_image_name(upgrade_image_type_en image_type)
{
	switch (image_type) {
		case IMAGE_TYPE_IS_UBOOT:
			return env_get("uboot_file");
		case IMAGE_TYPE_IS_DTB:
			return env_get("fdt_file");
		case IMAGE_TYPE_IS_KERNEL:
			return env_get("kernel_file");
		case IMAGE_TYPE_IS_ROOTFS:
			return env_get("rootfs_file");
		case IMAGE_TYPE_IS_ALL:
			return env_get("image");
		default:
			return NULL;
	}

	return NULL;
}

/**
 * @brief  设置下载镜像后存放的内存地址
 *
 * @return int 成功返回 0，失败返回 -1
 */
static int get_image_load_addr(upgrade_context_st *upgrade_context)
{
	if(NULL == upgrade_context) {
		return -EINVAL;
	}

	switch (upgrade_context->image_type) {
		case IMAGE_TYPE_IS_UBOOT:
		case IMAGE_TYPE_IS_KERNEL:
		case IMAGE_TYPE_IS_ROOTFS:
		case IMAGE_IS_ALL:
			upgrade_context->image_mem_addr = env_get("loadaddr");
			break;
		case IMAGE_TYPE_IS_DTB:
			upgrade_context->image_mem_addr = env_get("fdt_addr");
			break;
		default:
			return -EINVAL;
	}

	if(NULL == upgrade_context->image_mem_addr) {
		printf("The image is loaded in memory at a NULL address !!!\n");
		return -EINVAL;
	}

	return 0;
}

/**
 * @brief  从网络上获取镜像
 * @note   net_boot_file_name 是 net/net.c 中的全局变量，存储要下载的文件名
 * @param  file_size: 	输出，下载的文件的大小
 * @retval 成功返回传输的文件的大小，失败返回 -1
 */
static int fetch_image_from_net(upgrade_context_st *upgrade_context)
{
	int file_size = 0;

	copy_filename(net_boot_file_name, upgrade_context->image_name, sizeof(upgrade_context->image_name));

	file_size = net_loop(TFTPGET);
	if(file_size <= 0) return -1;
	upgrade_context->image_size = file_size;

	return 0;
}


static int mmc_update_image_uboot(int image_size, int)
{
	char cmd[256] = {0};
	int ret = 0;
	int bootsize = 0;

	/** 计算要烧写的大小，512 是一个块的大小 */
	bootsize = g_update_info.update_image_size / 512;

	switch(g_update_info.update_dev_type) {
		case BLK_DEV_SD:
			snprintf(cmd, sizeof(cmd), "mmc dev %d %d\n",
						g_update_info.sd_blk_info.mmc_dev_num, 
						g_update_info.sd_blk_info.bootloader_part_num0);
			ret = run_command(cmd, 0);	///< 切换到 EMMC 的分区 0，即 boot0 分区
			if(0 != ret) return ret;

			memset(cmd, 0, sizeof(cmd));
			/** 从 SD 分区 0 第 2 个块开始烧写新的 uboot，前 2 个块存放分区表 */
			snprintf(cmd, sizeof(cmd), "mmc write %x %d %x\n", 
						(unsigned int)g_update_info.load_addr, 
						g_update_info.sd_blk_info.bootloader_sector_offset,
						bootsize);
			ret = run_command(cmd, 0); ///< 烧写 u-boot.imx 到 EMMC 中
			if(0 != ret) return ret;
			break;
		case BLK_DEV_EMMC:
			snprintf(cmd, sizeof(cmd), "mmc dev %d %d\n",
						g_update_info.eMMC_blk_info.mmc_dev_num, 
						g_update_info.eMMC_blk_info.default_bootloader);
			ret = run_command(cmd, 0);	///< 切换到 EMMC 的 boot 分区 0，即 boot0 或 boot1 分区
			if(0 != ret) return ret;

			memset(cmd, 0, sizeof(cmd));
			/** 从 eMMC 分区 boot0 或 boot1 第 2 个块开始烧写新的 uboot，前 2 个块存放分区表 */
			snprintf(cmd, sizeof(cmd), "mmc write %x %d %x\n", 
						(unsigned int)g_update_info.load_addr, 
						g_update_info.eMMC_blk_info.bootloader_sector_offset,
						bootsize);
			ret = run_command(cmd, 0); ///< 烧写 u-boot.imx 到 EMMC 中
			if(0 != ret) return ret;

			/** eMMC 烧写完成后，需要执行 mmc partconf 命令来设置新的启动分区 */
			memset(cmd, 0, sizeof(cmd));
			snprintf(cmd, sizeof(cmd), "mmc partconf %d %d %d %d\n",
						g_update_info.eMMC_blk_info.mmc_dev_num,
						1,
						g_update_info.eMMC_blk_info.default_bootloader,
						g_update_info.eMMC_blk_info.default_bootloader);
			ret = run_command(cmd, 0); ///< 烧写 u-boot.imx 到 EMMC 中
			if(0 != ret) return ret;
			break;
		default:
			return -1;
	}

	return ret;
}

static int get_mmc_part_info(upgrade_media_type_en media_type,
									int dev_num, 
									mmc_part_info_st *part_info)
{
	struct disk_partition disk_part_info;
	int ret = -1;

	if(NULL == part_info || dev_num < 0) {
		printf("Invalid parameter !!! \n");
		return ret;
	}

	if (MEDIA_TYPE_IS_SD == media_type) {
		/** 需要注意的是 0 号分区通过此操作获取不到, 前 1024KB 为分区表项 */
		part_info->boot_part_num0 = 0;
		part_info->boot_sector_start = 0;
		part_info->default_boot_part = 0;
		// part_info->boot_sector_offset = 2;
	}
	else if (MEDIA_TYPE_IS_EMMC == media_type) {
		/** 需要注意的是 0 或 1 号分区通过此操作获取不到 */
		part_info->boot_part_num0 = 0;
		part_info->boot_part_num1 = 1;
		part_info->default_boot_part = 0;
		part_info->boot_sector_start = 0;
	 // part_info->boot_sector_offset = 2;
	}

	/** 获取 1 号分区信息 */
	ret = mmc_part_info_get(dev_num, 1, &disk_part_info);
	if(ret < 0) {
		printf("eMMC get part 1 info error.");
		return -1;
	}
	part_info->linux_sector_start = disk_part_info.start;
	part_info->linux_sector_offset = disk_part_info.size;
	part_info->linux_part_num = 1;

	/** 获取 2 号分区信息 */
	memset(&disk_part_info, 0, sizeof(disk_part_info));
	ret = mmc_part_info_get(dev_num, 2, &disk_part_info);
	if(ret < 0) {
		printf("eMMC get part 2 info error.");
		return -1;
	}
	part_info->rootfs_sector_start = disk_part_info.start;
	part_info->rootfs_sector_offset = disk_part_info.size;
	part_info->rootfs_part_num = 2;

	return 0;
}

static int get_mmc_dev_info(mmc_context_st *mmc_context, upgrade_media_type_en media_type)
{
	struct udevice *dev;
	char *mmc_type;

	if(NULL == mmc_context) {
		printf("mmc_context is NULL\n");
		return -1;
	}

	for (uclass_first_device(UCLASS_MMC, &dev);
	     dev;
	     uclass_next_device(&dev)) {
		struct mmc *m = mmc_get_mmc_dev(dev);

		if (m->has_init)
			mmc_type = IS_SD(m) ? "SD" : "eMMC";
		else
			mmc_type = NULL;

		if(strncmp(mmc_type, "SD", sizeof(mmc_type) == 0) 
			&& (MEDIA_TYPE_IS_SD == media_type)) {
			mmc_context->devnum = mmc_get_blk_desc(m)->devnum;
			mmc_context->mmc_size = m->capacity;
			get_mmc_part_info(media_type, mmc_context->devnum, &mmc_context->part_info);

			return 0;
		}

		if(strncmp(mmc_type, "eMMC", sizeof(mmc_type) == 0) 
			&& (MEDIA_TYPE_IS_EMMC == media_type)) {
			mmc_context->devnum = mmc_get_blk_desc(m)->devnum;
			mmc_context->mmc_size = m->capacity;
			get_mmc_part_info(media_type, mmc_context->devnum, &mmc_context->part_info);

			return 0;
		}
	}

	return 0;
}

/**
 * @brief 
 * 
 * @param upgrade_context 
 * @param mmc_context 
 * @return int 
 */
static int upgrade_uboot_to_mmc(upgrade_context_st *upgrade_context, 
								mmc_context_st *mmc_context)
{
	char cmd[256] = {0};
	int ret = -1;
	int bootsize = 0;

	if(NULL == upgrade_context || NULL == mmc_context) {
		printf("Invalid upgrade context or mmc context !!! \n");
		return ret;
	}

	/** 计算要烧写的大小，512 是一个块的大小 */
	bootsize = upgrade_context.image_size / 512;

	/** 切换到响应的分区 */
	snprintf(cmd, sizeof(cmd), "mmc dev %d %d\n",
				mmc_context->dev_num, 
				mmc_context->part_info.default_boot_part);
	ret = run_command(cmd, 0);	///< 切换到 mmc 设备的分区，eMMC 有 boot0 和 boot1
	if(0 != ret) return ret;

	/**  */
	memset(cmd, 0, sizeof(cmd));
	/** 从 MMC 的 boot 分区的第 2 个块开始烧写新的 uboot，前 2 个块存放分区表 */
	snprintf(cmd, sizeof(cmd), "mmc write %s 2 %x\n", 
				upgrade_context->image_mem_addr, 
				bootsize);
	ret = run_command(cmd, 0); ///< 烧写 u-boot 到 EMMC 中
	if(0 != ret) return ret;

	/** eMMC 烧写完成后，需要执行 mmc partconf 命令来设置新的启动分区 */
	if(MEDIA_TYPE_IS_EMMC == upgrade_context->media_type) {
		memset(cmd, 0, sizeof(cmd));
		snprintf(cmd, sizeof(cmd), "mmc partconf %d %d %d %d\n",
					mmc_context->dev_num,
					1,
					mmc_context->part_info.default_boot_part,
					mmc_context->part_info.default_boot_part);
		ret = run_command(cmd, 0); ///< 设置启动分区
		if(0 != ret) return ret;
	}

	return ret;
}

/**
 * @brief 
 * 
 * @param upgrade_context 
 * @param mmc_context 
 * @return int 
 */
static int upgrade_fdt_or_kernel_to_mmc(upgrade_context_st *upgrade_context, 
										mmc_context_st *mmc_context)
{
	int ret = -1;
	char cmd[256] = {0};

	if(NULL == upgrade_context || NULL == mmc_context) {
		printf("Invalid upgrade context or mmc context !!! \n");
		return ret;
	}

	memset(cmd, '\0', sizeof(cmd));
	/** 把内存上的数据存储到 fat 分区的一个文件里 */
	snprintf(cmd, sizeof(cmd), "fatwrite mmc %d:%d %s %s %x\n",
				mmc_context->dev_num, 
				mmc_context->part_info.linux_part_num, 
				upgrade_context->image_mem_addr, 
				upgrade_context->image_name
				upgrade_context->image_size);
	ret = run_command(cmd, 0);

	return ret;
}

/**
 * @brief 
 * 
 * @return int 
 */
static int upgrade_rootfs_to_mmc(upgrade_context_st *upgrade_context, 
								mmc_context_st *mmc_context)
{
	int ret = -1;
	char cmd[256] = {0};

	if(NULL == upgrade_context || NULL == mmc_context) {
		printf("Invalid upgrade context or mmc context !!! \n");
		return ret;
	}

	/** 文件系统的更新要先擦除掉整个分区 */
	/** 选中相应的 mmc 设备 */
	memset(cmd, '\0', sizeof(cmd));
	snprintf(cmd, sizeof(cmd), "mmc dev %d %d\n", mmc_context->dev_num);
	ret = run_command(cmd, 0);
	if(0 != ret) return ret;

	/** 擦除整个 ext4 分区 */
	memset(cmd, '\0', sizeof(cmd));
	snprintf(cmd, sizeof(cmd), "mmc erase %d %d\n", 
							mmc_context->part_info.rootfs_sector_start,
							mmc_context->part_info.rootfs_sector_offset);

	memset(cmd, '\0', sizeof(cmd));
	snprintf(cmd, sizeof(cmd), "ext4write mmc %d:%d %s %s %x\n",
				mmc_context->mmc_dev_num, 
				mmc_context->rootfs_part_num, 
				upgrade_context->image_mem_addr,
				upgrade_context->image_name,
				upgrade_context->image_size);
	/** 切换到文件系统分区 */
	ret = run_command(cmd, 0);
	if(ret != 0) return -1;

	return 0;
}

/**
 * @brief 
 * @note   更新的镜像为 sdcard.img
 * @return int 
 */
static int ugprade_sdcard_image_to_mmc(upgrade_context_st *upgrade_context, 
										mmc_context_st *mmc_context)
{
	int ret = -1;
	char cmd[256] = {0};

	if(NULL == upgrade_context || NULL == mmc_context) {
		printf("Invalid upgrade context or mmc context !!! \n");
		return ret;
	}

	/** 选中相应的 mmc 设备 */
	memset(cmd, '\0', sizeof(cmd));
	snprintf(cmd, sizeof(cmd), "mmc dev %d %d\n", mmc_context->dev_num);
	ret = run_command(cmd, 0);
	if(0 != ret) return ret;

	/** 擦除整个 mmc 设备 */
	memset(cmd, '\0', sizeof(cmd));
	snprintf(cmd, sizeof(cmd), "mmc erase %d %d\n", 
							mmc_context->part_info.rootfs_sector_start,
							mmc_context->part_info.rootfs_sector_offset);

	/** 刷写 sdcard.img 到 mmc 设备中 */

	return 0;
}

/**
 * @brief  烧写镜像
 * 
 * @return int 成功返回 0，失败返回 -1
 */
static int mmc_update_image(upgrade_context_st *context)
{
	mmc_context_st mmc_context = {0};

	if(NULL == context) {
		printf("Invalid upgrade context !!!\n");
		return -1;
	}

	/** 获取 mmc 设备的相关信息 */
	get_mmc_dev_info(&mmc_context, context->media_type);

	/** 执行升级 */
	switch (context->image_type) {
		case IMAGE_IS_UBOOT:	///< 更新 u-boot
			return upgrade_uboot_to_mmc(context, &mmc_context);
		case IMAGE_IS_DTB:		///< 更新设备树
		case IMAGE_IS_KERNEL:	///< 更新 kernel
			return mmc_update_image_fdt_or_kernel(context, &mmc_context);
		case IMAGE_IS_ROOTFS:	///< 更新文件系统
			return upgrade_rootfs_to_mmc(context, &mmc_context);
		case IMAGE_IS_ALL:		///< 更新打包的整体镜像
			return mmc_update_image_all();
		default:
			return -1;
	}

	return 0;
}

/**
 * @brief 
 * 
 * @param cmdtp 
 * @param flag 
 * @param argc 
 * @param argv 
 * @return int 
 */
static int do_upgrade_proc(upgrade_media_type_en media_type, int argc, char *const argv[])
{
    UPGRADE_IMAGE_TYPE_EN image_type = IMAGE_TYPE_IS_UNKOWN;
	char *image_name = NULL;
	upgrade_context_st upgrade_context = {0};

	if (argc != 3)
		return CMD_RET_USAGE;

	upgrade_context.media_type = media_type;

    upgrade_context.image_type = check_upgrade_image_type(argv[1]);
    if(IMAGE_TYPE_IS_UNKOWN == upgrade_context.image_type) {
        printf("Image type is unknown, please check your command line arguments !!!\n");
        return CMD_RET_FAILURE;
    }

	/** 如果命令行传入了镜像名，则直接使用传入的镜像 */
	if(NULL != argv[2]) {
		upgrade_context.image_name = argv[2];
	}
	else {
		/** 否则使用环境变量中的镜像名 */
		upgrade_context.image_name = check_upgrade_image_name(upgrade_context.image_type);
	}
	if (NULL == upgrade_context.image_name) {
		printf("Image type is unknown, please check your command line arguments !!!\n");
		return CMD_RET_FAILURE;
	}

	/** 设置 image_load_addr, 主要是设备树在内存中的地址与其他的不同 */
	if(get_image_load_addr(&upgrade_context) < 0) {
		printf("get image load addr failed !!!\n");
		return CMD_RET_FAILURE;
	}
	image_load_addr = simple_strtoul(upgrade_context.image_mem_addr, NULL, 16);

	/** 从网络上获取镜像 */
	if(fetch_image_from_net(&upgrade_context) < 0) {
		printf("Get image from network failed !!!\n");
		return CMD_RET_FAILURE;
	}

	/** 烧写镜像到指定地址, 若是 DRAM, 到这里基本就结束了 */
	switch (upgrade_context.media_type) {
		case MEDIA_TYPE_IS_DRAM:
			break;
		case MEDIA_TYPE_IS_SD:
		case MEDIA_TYPE_IS_EMMC:
			mmc_update_image(&upgrade_context);
			break;
		case MEDIA_TYPE_IS_SPI_FLASH:
			break;
		case MEDIA_TYPE_IS_QSPI_FLASH:
			break;
		case MEDIA_TYPE_IS_NAND_FLASH:
			break;
		default:
			break;
	}
    return 0;
}

static int do_upgrade_dram(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	do_upgrade_proc(int argc, char *const argv[]);
}

static int do_upgrade_sd(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc != 3)
		return CMD_RET_USAGE;

    return 0;
}

static int do_upgrade_emmc(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc != 3)
		return CMD_RET_USAGE;

    image_type = check_upgrade_image(argv[2]);
    if(IMAGE_TYPE_IS_UNKOWN == image_type) {
        printf("Image type is unknown, please check your command line arguments.\n");
        return CMD_RET_FAILURE;
    }

    return 0;
}

static int do_upgrade_spi_flash(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc != 3)
		return CMD_RET_USAGE;

    return 0;
}

static int do_upgrade_qspi_flash(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc != 3)
		return CMD_RET_USAGE;

    return 0;
}

static int do_upgrade_nand(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	if (argc != 3)
		return CMD_RET_USAGE;

    return 0;
}

static struct cmd_tbl cmd_upgrade_sub[] = {
	U_BOOT_CMD_MKENT(dram, 3, 0, do_upgrade_dram, "", ""),
	U_BOOT_CMD_MKENT(sd, 3, 0, do_upgrade_sd, "", ""),
	U_BOOT_CMD_MKENT(emmc, 3, 0, do_upgrade_emmc, "", ""),
	U_BOOT_CMD_MKENT(spi-flash, 3, 0, do_upgrade_spi_flash, "", ""),
	U_BOOT_CMD_MKENT(qspi-flash, 3, 0, do_upgrade_qspi_flash, "", ""),
	U_BOOT_CMD_MKENT(nand, 3, 0, do_upgrade_nand, "", ""),
};

static __maybe_unused void upgrade_reloc(void)
{
	static int relocated;

	if (!relocated) {
		fixup_cmdtable(cmd_upgrade_sub, ARRAY_SIZE(cmd_upgrade_sub));
		relocated = 1;
	};
}

static int do_upgrade(struct cmd_tbl *cmdtp, int flag, int argc, char *const argv[])
{
	struct cmd_tbl *c;

#ifdef CONFIG_NEEDS_MANUAL_RELOC
	upgrade_reloc();
#endif

	if (argc < 4)
		return CMD_RET_USAGE;

	/* Strip off leading 'upgrade' command argument */
	argc--;
	argv++;

	c = find_cmd_tbl(argv[0], &cmd_upgrade_sub[0], ARRAY_SIZE(cmd_upgrade_sub));

	if (c)
		return c->cmd(cmdtp, flag, argc, argv);
	else
		return CMD_RET_USAGE;
}

U_BOOT_CMD(
	upgrade, 4, 0, do_upgrade,
	"upgrade image to storage device",
    "upgrade dram [type] [name] - upgrade image to dram\n"
    "upgrade sd [type] [name] - upgrade image to emmc\n"
    "upgrade emmc [type] [name] - upgrade image to emmc\n"
    "upgrade spi-flash [type] [name] - upgrade image to nor-flash\n"
    "upgrade qspi-flash [type] [name] - upgrade image to qspi-flash\n"
    "upgrade nand [type] [name] - upgrade image to nand\n"
    "   arguments: \n"
    "       - type  : image type, It can only be uboot, kernel, dts, rootfs, and image.\n"
    "       - name  : image name.\n"
)